import re
import json
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QVBoxLayout, QHBoxLayout, QWidget,
    QLabel, QPushButton, QTextEdit, QFileDialog, QComboBox,
    QGroupBox, QLineEdit, QMessageBox, QCheckBox, QGridLayout
)
from PyQt6.QtCore import Qt


class RegexExtractorApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("文本正则提取工具")
        self.setGeometry(100, 100, 900, 700)
        
        self.init_ui()
        self.setup_regex_patterns()
        
    def init_ui(self):
        # 主布局
        main_widget = QWidget()
        main_layout = QVBoxLayout()
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)
        
        # 文件选择部分
        file_group = QGroupBox("文件选择")
        file_layout = QHBoxLayout()
        
        self.file_path_edit = QLineEdit()
        self.file_path_edit.setPlaceholderText("选择文件...")
        file_layout.addWidget(self.file_path_edit)
        
        browse_btn = QPushButton("浏览...")
        browse_btn.clicked.connect(self.browse_file)
        file_layout.addWidget(browse_btn)
        
        file_group.setLayout(file_layout)
        main_layout.addWidget(file_group)
        
        # 正则表达式选择部分
        regex_group = QGroupBox("正则表达式选项")
        regex_layout = QVBoxLayout()
        
        # 预设正则选择
        self.regex_preset_combo = QComboBox()
        self.regex_preset_combo.addItem("选择预设正则表达式...")
        regex_layout.addWidget(self.regex_preset_combo)
        
        # 自定义正则
        custom_regex_layout = QHBoxLayout()
        custom_regex_layout.addWidget(QLabel("自定义正则:"))
        
        self.custom_regex_edit = QLineEdit()
        self.custom_regex_edit.setPlaceholderText("输入自定义正则表达式")
        custom_regex_layout.addWidget(self.custom_regex_edit)
        
        regex_layout.addLayout(custom_regex_layout)
        
        # 指定文件后缀部分
        self.ext_group = QGroupBox("指定文件后缀")
        ext_layout = QGridLayout()
        
        # 常用后缀复选框
        self.common_extensions = [
            ('txt', '文本文件'), ('pdf', 'PDF文件'), 
            ('docx', 'Word文档'), ('xlsx', 'Excel文件'),
            ('pptx', 'PPT文件'), ('jpg', 'JPEG图片'),
            ('png', 'PNG图片'), ('mp3', 'MP3音频'),
            ('mp4', 'MP4视频'), ('zip', 'ZIP压缩包')
        ]
        
        self.ext_checkboxes = []
        for i, (ext, desc) in enumerate(self.common_extensions):
            cb = QCheckBox(f".{ext} ({desc})")
            self.ext_checkboxes.append((ext, cb))
            ext_layout.addWidget(cb, i // 3, i % 3)
        
        # 自定义后缀输入
        custom_ext_layout = QHBoxLayout()
        custom_ext_layout.addWidget(QLabel("自定义后缀:"))
        
        self.custom_ext_edit = QLineEdit()
        self.custom_ext_edit.setPlaceholderText("多个后缀用空格或|分隔，如: php asp jsp")
        custom_ext_layout.addWidget(self.custom_ext_edit)
        
        ext_layout.addLayout(custom_ext_layout, len(self.common_extensions) // 3 + 1, 0, 1, 3)
        self.ext_group.setLayout(ext_layout)
        regex_layout.addWidget(self.ext_group)
        
        # 上下文提取选项
        self.context_group = QGroupBox("上下文提取选项")
        context_layout = QVBoxLayout()
        
        # 上下文提取启用复选框
        self.enable_context_check = QCheckBox("启用上下文提取")
        self.enable_context_check.stateChanged.connect(self.toggle_context_mode)
        context_layout.addWidget(self.enable_context_check)
        
        # 保留前后文复选框
        self.keep_context_check = QCheckBox("保留前后文本")
        self.keep_context_check.setChecked(False)
        self.keep_context_check.setEnabled(False)
        context_layout.addWidget(self.keep_context_check)
        
        # 前后固定文本
        context_text_layout = QHBoxLayout()
        
        context_text_layout.addWidget(QLabel("前文本:"))
        self.before_text_edit = QLineEdit()
        self.before_text_edit.setEnabled(False)
        context_text_layout.addWidget(self.before_text_edit)
        
        context_text_layout.addWidget(QLabel("后文本:"))
        self.after_text_edit = QLineEdit()
        self.after_text_edit.setEnabled(False)
        context_text_layout.addWidget(self.after_text_edit)
        
        context_layout.addLayout(context_text_layout)
        self.context_group.setLayout(context_layout)
        regex_layout.addWidget(self.context_group)
        
        # 其他选项
        options_layout = QHBoxLayout()
        
        self.ignore_case_check = QCheckBox("忽略大小写")
        options_layout.addWidget(self.ignore_case_check)
        
        self.multiline_check = QCheckBox("多行模式")
        options_layout.addWidget(self.multiline_check)
        
        self.dotall_check = QCheckBox("点号匹配换行")
        options_layout.addWidget(self.dotall_check)
        
        regex_layout.addLayout(options_layout)
        
        regex_group.setLayout(regex_layout)
        main_layout.addWidget(regex_group)
        
        # 操作按钮
        button_layout = QHBoxLayout()
        
        extract_btn = QPushButton("提取内容")
        extract_btn.clicked.connect(self.extract_content)
        button_layout.addWidget(extract_btn)
        
        save_btn = QPushButton("保存结果")
        save_btn.clicked.connect(self.save_results)
        button_layout.addWidget(save_btn)
        
        clear_btn = QPushButton("清空")
        clear_btn.clicked.connect(self.clear_all)
        button_layout.addWidget(clear_btn)
        
        gen_ext_regex_btn = QPushButton("生成后缀正则")
        gen_ext_regex_btn.clicked.connect(self.generate_ext_regex)
        button_layout.addWidget(gen_ext_regex_btn)
        
        main_layout.addLayout(button_layout)
        
        # 结果显示
        result_group = QGroupBox("提取结果")
        result_layout = QVBoxLayout()
        
        self.result_edit = QTextEdit()
        self.result_edit.setReadOnly(False)
        result_layout.addWidget(self.result_edit)
        
        result_group.setLayout(result_layout)
        main_layout.addWidget(result_group)
    
    def toggle_context_mode(self, state):
        """切换上下文提取模式"""
        enable = state == Qt.CheckState.Checked.value
        
        # 启用/禁用上下文相关控件
        self.keep_context_check.setEnabled(enable)
        self.before_text_edit.setEnabled(enable)
        self.after_text_edit.setEnabled(enable)
        
        # 如果启用上下文提取，则禁用预设正则和文件后缀选项
        self.regex_preset_combo.setEnabled(not enable)
        self.custom_regex_edit.setEnabled(not enable)
        self.ext_group.setEnabled(not enable)
        
        # 更新UI状态
        if enable:
            self.context_group.setStyleSheet("QGroupBox { border: 2px solid blue; }")
            self.regex_preset_combo.setCurrentIndex(0)
            self.custom_regex_edit.clear()
        else:
            self.context_group.setStyleSheet("")
    
    def setup_regex_patterns(self):
        # 预设正则表达式
        self.regex_presets = [
            ("URL网址 (完整)", r'(?:(?:https?|ftp|file)://|www\.|ftp\.)(?:\([-A-Z0-9+&@#/%=~_|$?!:,.]*\)|[-A-Z0-9+&@#/%=~_|$?!:,.])*(?:\([-A-Z0-9+&@#/%=~_|$?!:,.]*\)|[A-Z0-9+&@#/%=~_|$])', re.IGNORECASE),
            ("URL网址 (简版)", r'https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+[/\w .-]*/?', re.IGNORECASE),
            ("图片链接", r'(?:(?:https?|ftp|file)://|www\.|ftp\.)(?:\([-A-Z0-9+&@#/%=~_|$?!:,.]*\)|[-A-Z0-9+&@#/%=~_|$?!:,.])*\.(?:jpg|jpeg|png|gif|bmp|webp|svg)(?:\?[^"\'\s]*)?', re.IGNORECASE),
            ("音乐链接", r'(?:(?:https?|ftp|file)://|www\.|ftp\.)(?:\([-A-Z0-9+&@#/%=~_|$?!:,.]*\)|[-A-Z0-9+&@#/%=~_|$?!:,.])*\.(?:mp3|wav|flac|aac|ogg|wma)(?:\?[^"\'\s]*)?', re.IGNORECASE),
            ("视频链接", r'(?:(?:https?|ftp|file)://|www\.|ftp\.)(?:\([-A-Z0-9+&@#/%=~_|$?!:,.]*\)|[-A-Z0-9+&@#/%=~_|$?!:,.])*\.(?:mp4|avi|mov|flv|wmv|mkv|webm)(?:\?[^"\'\s]*)?', re.IGNORECASE),
            ("网页链接 (仅http/https)", r'https?://(?:[-\w.]|(?:%[\da-fA-F]{2}))+(?:[/\w .-]*/?)?(?:\?[^"\'\s]*)?(?:#[^"\'\s]*)?', re.IGNORECASE),
            ("身份证号码", r'[1-9]\d{5}(18|19|20)\d{2}(0[1-9]|1[0-2])(0[1-9]|[12]\d|3[01])\d{3}[\dXx]'),
            ("手机号码", r'(?<!\d)(1[3-9]\d{9})(?!\d)'),
            ("固定电话号码", r'(?<!\d)(\d{3,4}-\d{7,8})(?!\d)'),
            ("日期 (YYYY-MM-DD)", r'\d{4}-\d{1,2}-\d{1,2}'),
            ("日期 (YYYY/MM/DD)", r'\d{4}/\d{1,2}/\d{1,2}'),
            ("邮箱地址", r'[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}'),
            ("IP地址", r'\b(?:\d{1,3}\.){3}\d{1,3}\b'),
            ("中文", r'[\u4e00-\u9fa5]+'),
            ("HTML标签", r'<[^>]+>'),
            ("指定后缀文件", r'\b\w+\.(txt|pdf|docx?|xlsx?|pptx?|zip|rar|7z)\b')
        ]
        
        for name, pattern, *flags in self.regex_presets:
            flag = flags[0] if flags else 0
            self.regex_preset_combo.addItem(name)
        
        self.regex_preset_combo.currentIndexChanged.connect(self.on_regex_preset_changed)
    
    def generate_ext_regex(self):
        """生成文件后缀正则并显示在自定义正则输入框"""
        if self.enable_context_check.isChecked():
            QMessageBox.warning(self, "警告", "当前处于上下文提取模式，请先禁用上下文提取")
            return
            
        ext_regex = self.get_file_extensions_regex()
        self.custom_regex_edit.setText(ext_regex)
        QMessageBox.information(self, "提示", f"已生成后缀正则:\n{ext_regex}")
    
    def get_file_extensions_regex(self):
        """获取选中的文件后缀正则表达式"""
        selected_exts = []
        
        # 获取选中的常见后缀
        for ext, cb in self.ext_checkboxes:
            if cb.isChecked():
                selected_exts.append(ext)
        
        # 获取自定义后缀
        custom_exts = self.custom_ext_edit.text().strip()
        if custom_exts:
            # 支持用空格、逗号、竖线分隔
            for ext in re.split(r'[\s,|]+', custom_exts):
                ext = ext.strip().lstrip('.')  # 去除可能的点和空白
                if ext:
                    selected_exts.append(ext)
        
        if not selected_exts:
            return r'\b\w+\.\w+\b'  # 默认匹配所有文件后缀
        
        # 去重并排序
        selected_exts = sorted(list(set(selected_exts)))
        
        # 构建正则表达式，支持忽略大小写
        ext_pattern = '|'.join(re.escape(ext) for ext in selected_exts)
        return fr'\b\w+\.({ext_pattern})\b'
    
    def on_regex_preset_changed(self, index):
        if index > 0 and not self.enable_context_check.isChecked():
            pattern_info = self.regex_presets[index-1]
            pattern = pattern_info[1]
            
            # 如果是"指定后缀文件"，使用动态生成的正则
            if pattern_info[0] == "指定后缀文件":
                pattern = self.get_file_extensions_regex()
                self.ext_group.setStyleSheet("QGroupBox { border: 2px solid blue; }")
            else:
                self.ext_group.setStyleSheet("")
            
            self.custom_regex_edit.setText(pattern)
            
            # 自动设置标志
            if len(pattern_info) > 2 and pattern_info[2] & re.IGNORECASE:
                self.ignore_case_check.setChecked(True)
            else:
                self.ignore_case_check.setChecked(False)
    
    def browse_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择文件", "", "文本文件 (*.txt);;JSON文件 (*.json);;所有文件 (*.*)"
        )
        if file_path:
            self.file_path_edit.setText(file_path)
    
    def get_file_content(self):
        file_path = self.file_path_edit.text()
        if not file_path:
            QMessageBox.warning(self, "警告", "请先选择文件!")
            return None
        
        try:
            if file_path.lower().endswith('.json'):
                with open(file_path, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    return json.dumps(data, ensure_ascii=False, indent=2)
            else:
                with open(file_path, 'r', encoding='utf-8') as f:
                    return f.read()
        except Exception as e:
            QMessageBox.critical(self, "错误", f"读取文件失败: {str(e)}")
            return None
    
    def extract_content(self):
        content = self.get_file_content()
        if content is None:
            return
        
        # 设置正则标志
        flags = 0
        if self.ignore_case_check.isChecked():
            flags |= re.IGNORECASE
        if self.multiline_check.isChecked():
            flags |= re.MULTILINE
        if self.dotall_check.isChecked():
            flags |= re.DOTALL
        
        try:
            if self.enable_context_check.isChecked():
                # 上下文提取模式
                before_text = self.before_text_edit.text().strip()
                after_text = self.after_text_edit.text().strip()
                
                if not before_text and not after_text:
                    QMessageBox.warning(self, "警告", "请至少输入前文本或后文本!")
                    return
                
                if self.keep_context_check.isChecked():
                    # 保留前后文本的模式
                    pattern_str = ''
                    
                    if before_text:
                        pattern_str += f'({re.escape(before_text)})'
                    
                    # 匹配任意内容
                    pattern_str += r'(.*?)'
                    
                    if after_text:
                        pattern_str += f'({re.escape(after_text)})'
                    
                    pattern = re.compile(pattern_str, flags)
                    matches = pattern.finditer(content)
                    
                    # 组合前后文和匹配内容
                    results = []
                    for match in matches:
                        parts = []
                        if before_text:
                            parts.append(match.group(1))
                        parts.append(match.group(2))
                        if after_text:
                            parts.append(match.group(3))
                        results.append(''.join(parts))
                else:
                    # 不保留前后文本的模式
                    pattern_str = ''
                    
                    if before_text:
                        pattern_str += f'(?<={re.escape(before_text)})'
                    
                    # 匹配任意内容
                    pattern_str += r'(.*?)'
                    
                    if after_text:
                        pattern_str += f'(?={re.escape(after_text)})'
                    
                    pattern = re.compile(pattern_str, flags)
                    matches = pattern.finditer(content)
                    
                    # 只提取匹配的内容部分
                    results = [match.group(1) for match in matches if match.group(1).strip()]
            else:
                # 普通正则提取模式
                if self.regex_preset_combo.currentIndex() > 0:
                    selected_preset = self.regex_presets[self.regex_preset_combo.currentIndex()-1][0]
                    if selected_preset == "指定后缀文件":
                        self.custom_regex_edit.setText(self.get_file_extensions_regex())
                
                regex_pattern = self.custom_regex_edit.text().strip()
                if not regex_pattern:
                    QMessageBox.warning(self, "警告", "请输入正则表达式!")
                    return
                
                pattern = re.compile(regex_pattern, flags)
                matches = pattern.finditer(content)
                results = [match.group() for match in matches if match.group().strip()]
            
            if not results:
                self.result_edit.setPlainText("没有找到匹配的内容。")
                return
            
            # 显示结果
            result_text = "\n\n".join(results)
            self.result_edit.setPlainText(f"共找到 {len(results)} 个匹配项:\n\n{result_text}")
            
        except re.error as e:
            QMessageBox.critical(self, "正则表达式错误", f"无效的正则表达式: {str(e)}")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"提取内容时出错: {str(e)}")
    
    def save_results(self):
        if not self.result_edit.toPlainText():
            QMessageBox.warning(self, "警告", "没有可保存的结果!")
            return
        
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存结果", "", "文本文件 (*.txt);;JSON文件 (*.json);;所有文件 (*.*)"
        )
        if file_path:
            try:
                with open(file_path, 'w', encoding='utf-8') as f:
                    if file_path.lower().endswith('.json'):
                        results = [r for r in self.result_edit.toPlainText().split('\n\n') if r.strip()]
                        json.dump(results, f, ensure_ascii=False, indent=2)
                    else:
                        f.write(self.result_edit.toPlainText())
                QMessageBox.information(self, "成功", "结果保存成功!")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"保存文件失败: {str(e)}")
    
    def clear_all(self):
        self.result_edit.clear()


if __name__ == "__main__":
    app = QApplication([])
    window = RegexExtractorApp()
    window.show()
    app.exec()